// SPDX-FileCopyrightText: © 2023 Tenstorrent AI ULC
//
// SPDX-License-Identifier: Apache-2.0
#include <core_descriptor.hpp>
#include <device.hpp>
#include <device_pool.hpp>
#include <dispatch_core_common.hpp>
#include <host_api.hpp>
#include <profiler.hpp>
#include <mesh_workload.hpp>
#include <mesh_command_queue.hpp>
#include <tt_metal.hpp>
#include <tt_metal_profiler.hpp>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <exception>
#include <filesystem>
#include <limits>
#include <map>
#include <optional>
#include <ostream>
#include <set>
#include <string>
#include <thread>
#include <tuple>
#include <unordered_map>
#include <utility>
#include <variant>
#include <vector>

#include <tt_stl/assert.hpp>
#include "buffer.hpp"
#include "core_coord.hpp"
#include "data_types.hpp"
#include "hal_types.hpp"
#include "hostdevcommon/profiler_common.h"
#include "impl/context/metal_context.hpp"
#include "kernel_types.hpp"
#include "llrt.hpp"
#include "llrt/hal.hpp"
#include <tt-logger/tt-logger.hpp>
#include "metal_soc_descriptor.h"
#include "profiler_optional_metadata.hpp"
#include "profiler_paths.hpp"
#include "profiler_state.hpp"
#include "profiler_types.hpp"
#include "profiler_state_manager.hpp"
#include "tt-metalium/program.hpp"
#include <tt-metalium/device_pool.hpp>
#include "rtoptions.hpp"
#include "tracy/Tracy.hpp"
#include "tracy/TracyTTDevice.hpp"
#include <tt-metalium/distributed.hpp>
#include <umd/device/types/core_coordinates.hpp>
#include <umd/device/types/xy_pair.hpp>

namespace tt {

namespace tt_metal {

namespace detail {

void setControlBuffer(IDevice* device, std::vector<uint32_t>& control_buffer) {
#if defined(TRACY_ENABLE)
    const chip_id_t device_id = device->id();
    const metal_SocDescriptor& soc_d = tt::tt_metal::MetalContext::instance().get_cluster().get_soc_desc(device_id);

    control_buffer[kernel_profiler::CORE_COUNT_PER_DRAM] = soc_d.profiler_ceiled_core_count_perf_dram_bank;
    for (auto core :
         tt::tt_metal::MetalContext::instance().get_cluster().get_virtual_routing_to_profiler_flat_id(device_id)) {
        const CoreCoord curr_core = core.first;

        control_buffer[kernel_profiler::FLAT_ID] = core.second;

        writeToCoreControlBuffer(device, curr_core, control_buffer);
    }
#endif
}

void syncDeviceHost(IDevice* device, CoreCoord logical_core, bool doHeader) {
    ZoneScopedC(tracy::Color::Tomato3);
    if (!tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_sync_enabled()) {
        return;
    }
    auto device_id = device->id();
    auto core = device->worker_core_from_logical_core(logical_core);

    const metal_SocDescriptor& soc_desc = tt::tt_metal::MetalContext::instance().get_cluster().get_soc_desc(device_id);
    auto phys_core = soc_desc.translate_coord_to(core, CoordSystem::TRANSLATED, CoordSystem::NOC0);

    const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
        tt::tt_metal::MetalContext::instance().profiler_state_manager();

    profiler_state_manager->device_host_time_pair.emplace(device_id, (std::vector<std::pair<uint64_t, uint64_t>>){});
    profiler_state_manager->smallest_host_time.emplace(device_id, 0);

    constexpr uint16_t sampleCount = 249;
    // TODO(MO): Always recreate a new program until subdevice
    // allows using the first program generated by default manager
    tt_metal::Program sync_program;

    std::map<std::string, std::string> kernel_defines = {
        {"SAMPLE_COUNT", std::to_string(sampleCount)},
    };

    tt_metal::CreateKernel(
        sync_program,
        "tt_metal/tools/profiler/sync/sync_kernel.cpp",
        logical_core,
        tt_metal::DataMovementConfig{
            .processor = tt_metal::DataMovementProcessor::RISCV_0,
            .noc = tt_metal::NOC::RISCV_0_default,
            .defines = kernel_defines});

    // Using MeshDevice APIs if the current device is managed by MeshDevice
    tt_metal::detail::LaunchProgram(
        device, sync_program, false /* wait_until_cores_done */, /* force_slow_dispatch */ true);

    std::filesystem::path output_dir = std::filesystem::path(get_profiler_logs_dir());
    std::filesystem::path log_path = output_dir / "sync_device_info.csv";
    std::ofstream log_file;

    constexpr int millisecond_wait = 10;

    const double tracyToSecRatio = TracyGetTimerMul();
    const int64_t tracyBaseTime = TracyGetBaseTime();
    const int64_t hostStartTime = TracyGetCpuTime();
    std::vector<int64_t> writeTimes(sampleCount);

    const auto& hal = MetalContext::instance().hal();
    HalProgrammableCoreType core_type = device->get_programmable_core_type(core);
    auto dev_msgs_factory = hal.get_dev_msgs_factory(core_type);
    DeviceAddr profiler_msg_addr = hal.get_dev_addr(core_type, HalL1MemAddrType::PROFILER);
    DeviceAddr control_vector_addr = profiler_msg_addr + dev_msgs_factory.offset_of<dev_msgs::profiler_msg_t>(
                                                             dev_msgs::profiler_msg_t::Field::control_vector);
    DeviceAddr control_addr = control_vector_addr + (kernel_profiler::FW_RESET_L * sizeof(uint32_t));
    for (int i = 0; i < sampleCount; i++) {
        ZoneScopedC(tracy::Color::Tomato2);
        std::this_thread::sleep_for(std::chrono::milliseconds(millisecond_wait));
        int64_t writeStart = TracyGetCpuTime();
        uint32_t sinceStart = writeStart - hostStartTime;

        tt::tt_metal::MetalContext::instance().get_cluster().write_reg(
            &sinceStart, tt_cxy_pair(device_id, core), control_addr);
        writeTimes[i] = (TracyGetCpuTime() - writeStart);
    }
    tt_metal::detail::WaitProgramDone(device, sync_program, false);
    std::vector<CoreCoord> cores = {core};
    profiler_state_manager->device_profiler_map.at(device_id).readResults(
        device, cores, ProfilerReadState::NORMAL, ProfilerDataBufferSource::L1);
    profiler_state_manager->device_profiler_map.at(device_id).processResults(
        device, cores, ProfilerReadState::NORMAL, ProfilerDataBufferSource::L1);

    log_info(tt::LogMetal, "SYNC PROGRAM FINISH IS DONE ON {}", device_id);
    if ((profiler_state_manager->smallest_host_time.at(device_id) == 0) ||
        (profiler_state_manager->smallest_host_time.at(device_id) > hostStartTime)) {
        profiler_state_manager->smallest_host_time.at(device_id) = hostStartTime;
    }

    uint64_t addr = profiler_msg_addr +
                    dev_msgs_factory.offset_of<dev_msgs::profiler_msg_t>(dev_msgs::profiler_msg_t::Field::buffer) +
                    (kernel_profiler::CUSTOM_MARKERS * sizeof(uint32_t));

    std::vector<std::uint32_t> sync_times = tt::tt_metal::MetalContext::instance().get_cluster().read_core(
        device_id, core, addr, (sampleCount + 1) * 2 * sizeof(uint32_t));

    uint32_t preDeviceTime = 0;
    uint32_t preHostTime = 0;
    bool firstSample = true;

    uint32_t deviceStartTime_H = sync_times[0] & 0xFFF;
    uint32_t deviceStartTime_L = sync_times[1];
    preDeviceTime = deviceStartTime_L;

    uint32_t hostStartTime_H = 0;

    for (int i = 2; i < 2 * (sampleCount + 1); i += 2) {
        uint32_t deviceTime = sync_times[i];
        if (deviceTime < preDeviceTime) {
            deviceStartTime_H++;
        }
        preDeviceTime = deviceTime;
        uint64_t deviceTimeLarge = (uint64_t(deviceStartTime_H) << 32) | deviceTime;

        uint32_t hostTime = sync_times[i + 1] + writeTimes[(i / 2) - 1];
        if (hostTime < preHostTime) {
            hostStartTime_H++;
        }
        preHostTime = hostTime;
        uint64_t hostTimeLarge = hostStartTime - profiler_state_manager->smallest_host_time.at(device_id) +
                                 ((uint64_t(hostStartTime_H) << 32) | hostTime);

        profiler_state_manager->device_host_time_pair.at(device_id).push_back(
            std::pair<uint64_t, uint64_t>{deviceTimeLarge, hostTimeLarge});

        if (firstSample) {
            firstSample = false;
        }
    }

    double hostSum = 0;
    double deviceSum = 0;
    double hostSquaredSum = 0;
    double hostDeviceProductSum = 0;

    for (auto& deviceHostTime : profiler_state_manager->device_host_time_pair.at(device_id)) {
        double deviceTime = deviceHostTime.first;
        double hostTime = deviceHostTime.second;

        deviceSum += deviceTime;
        hostSum += hostTime;
        hostSquaredSum += (hostTime * hostTime);
        hostDeviceProductSum += (hostTime * deviceTime);
    }

    uint16_t accumulateSampleCount = profiler_state_manager->device_host_time_pair.at(device_id).size();

    double frequencyFit = (hostDeviceProductSum * accumulateSampleCount - hostSum * deviceSum) /
                          ((hostSquaredSum * accumulateSampleCount - hostSum * hostSum) * tracyToSecRatio);

    double delay = (deviceSum - frequencyFit * hostSum * tracyToSecRatio) / accumulateSampleCount;

    if (doHeader) {
        log_file.open(log_path);
        log_file << fmt::format(
                        "device id,core_x, "
                        "core_y,device,host_tracy,host_real,write_overhead,host_start,delay,frequency,tracy_ratio,"
                        "tracy_base_time,device_frequency_ratio,device_shift")
                 << std::endl;
    } else {
        log_file.open(log_path, std::ios_base::app);
    }

    int init = profiler_state_manager->device_host_time_pair.at(device_id).size() - sampleCount;
    for (int i = init; i < profiler_state_manager->device_host_time_pair.at(device_id).size(); i++) {
        log_file << fmt::format(
                        "{:5},{:5},{:5},{:20},{:20},{:20.2f},{:20},{:20},{:20.2f},{:20.15f},{:20.15f},{:20},1.0,0",
                        device_id,
                        phys_core.x,
                        phys_core.y,
                        profiler_state_manager->device_host_time_pair.at(device_id)[i].first,
                        profiler_state_manager->device_host_time_pair.at(device_id)[i].second,
                        (double)profiler_state_manager->device_host_time_pair.at(device_id)[i].second * tracyToSecRatio,
                        writeTimes[i - init],
                        profiler_state_manager->smallest_host_time.at(device_id),
                        delay,
                        frequencyFit,
                        tracyToSecRatio,
                        tracyBaseTime)
                 << std::endl;
    }
    log_file.close();
    log_info(
        tt::LogMetal,
        "Host sync data for device: {}, cpu_start:{}, delay:{}, freq:{} Hz",
        device_id,
        profiler_state_manager->smallest_host_time.at(device_id),
        delay,
        frequencyFit);

    double host_timestamp = hostStartTime;
    double device_timestamp = delay + ((host_timestamp - profiler_state_manager->smallest_host_time.at(device_id)) *
                                       frequencyFit * tracyToSecRatio);
    // disable linting here; slicing is __intended__
    // NOLINTBEGIN
    profiler_state_manager->device_profiler_map.at(device_id).device_core_sync_info.insert_or_assign(
        CoreCoord(phys_core), SyncInfo(host_timestamp, device_timestamp, frequencyFit));
    // NOLINTEND
}

void setShift(int device_id, int64_t shift, double scale, const SyncInfo& root_sync_info) {
    if (std::isnan(scale)) {
        return;
    }
    log_info(tt::LogMetal, "Device sync data for device: {}, delay: {} ns, freq scale: {}", device_id, shift, scale);
    if (tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_mid_run_dump()) {
        log_warning(
            tt::LogMetal,
            "Note that tracy mid-run data dumping is enabled. This means device-device sync is not as accurate. Please "
            "do not use tracy mid-run data dumping for sensitive device-device event analysis.");
    }

    const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
        tt::tt_metal::MetalContext::instance().profiler_state_manager();

    auto device_profiler_it = profiler_state_manager->device_profiler_map.find(device_id);
    if (device_profiler_it != profiler_state_manager->device_profiler_map.end()) {
        device_profiler_it->second.freq_scale = scale;
        device_profiler_it->second.shift = shift;
        device_profiler_it->second.setSyncInfo(root_sync_info);

        std::filesystem::path output_dir = std::filesystem::path(get_profiler_logs_dir());
        std::filesystem::path log_path = output_dir / "sync_device_info.csv";
        std::ofstream log_file;
        log_file.open(log_path, std::ios_base::app);
        log_file << fmt::format("{:5},,,,,,,,,,,,{:20.15f},{:20}", device_id, scale, shift) << std::endl;
        log_file.close();
    }
}

void peekDeviceData(IDevice* device, std::vector<CoreCoord>& worker_cores) {
    ZoneScoped;
    auto device_id = device->id();
    std::string zoneName = fmt::format("peek {}", device_id);
    ZoneName(zoneName.c_str(), zoneName.size());
    const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
        tt::tt_metal::MetalContext::instance().profiler_state_manager();
    const auto& device_profiler_it = profiler_state_manager->device_profiler_map.find(device_id);
    if (device_profiler_it != profiler_state_manager->device_profiler_map.end()) {
        DeviceProfiler& device_profiler = device_profiler_it->second;
        device_profiler.device_sync_new_markers.clear();
        device_profiler.readResults(device, worker_cores, ProfilerReadState::NORMAL, ProfilerDataBufferSource::L1);
        device_profiler.processResults(device, worker_cores, ProfilerReadState::NORMAL, ProfilerDataBufferSource::L1);
        for (const auto& [core, risc_map] : device_profiler.device_markers_per_core_risc_map) {
            for (const auto& [risc, device_markers] : risc_map) {
                for (const tracy::TTDeviceMarker& marker : device_markers) {
                    const tracy::MarkerDetails marker_details = device_profiler.getMarkerDetails(marker.marker_id);
                    if (marker_details.marker_name_keyword_flags[static_cast<uint16_t>(
                            tracy::MarkerDetails::MarkerNameKeyword::SYNC_ZONE)]) {
                        ZoneScopedN("Adding_device_sync_marker");
                        auto ret = device_profiler.device_sync_markers.insert(marker);
                        if (ret.second) {
                            device_profiler.device_sync_new_markers.insert(marker);
                        }
                    }
                }
            }
        }
    }
}

void syncDeviceDevice(chip_id_t device_id_sender, chip_id_t device_id_receiver) {
    ZoneScopedC(tracy::Color::Tomato4);
    std::string zoneName = fmt::format("sync_device_device_{}->{}", device_id_sender, device_id_receiver);
    ZoneName(zoneName.c_str(), zoneName.size());
    if (!tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_sync_enabled()) {
        return;
    }

    IDevice* device_sender = nullptr;
    IDevice* device_receiver = nullptr;

    if (tt::DevicePool::instance().is_device_active(device_id_receiver)) {
        device_receiver = tt::DevicePool::instance().get_active_device(device_id_receiver);
    }

    if (tt::DevicePool::instance().is_device_active(device_id_sender)) {
        device_sender = tt::DevicePool::instance().get_active_device(device_id_sender);
    }

    if (device_sender != nullptr and device_receiver != nullptr) {
        constexpr std::uint16_t sample_count = 240;
        constexpr std::uint16_t sample_size = 16;
        constexpr std::uint16_t channel_count = 1;

        const auto& active_eth_cores = device_sender->get_active_ethernet_cores(false);
        auto eth_sender_core_iter = active_eth_cores.begin();
        tt_xy_pair eth_receiver_core;
        tt_xy_pair eth_sender_core;

        chip_id_t device_id_receiver_curr = std::numeric_limits<chip_id_t>::max();
        while ((device_id_receiver != device_id_receiver_curr) and (eth_sender_core_iter != active_eth_cores.end())) {
            eth_sender_core = *eth_sender_core_iter;
            if (not tt::tt_metal::MetalContext::instance().get_cluster().is_ethernet_link_up(
                    device_sender->id(), eth_sender_core)) {
                eth_sender_core_iter++;
                continue;
            }
            std::tie(device_id_receiver_curr, eth_receiver_core) =
                device_sender->get_connected_ethernet_core(eth_sender_core);
            eth_sender_core_iter++;
        }

        if (device_id_receiver != device_id_receiver_curr) {
            log_warning(
                tt::LogMetal,
                "No eth connection could be found between device {} and {}",
                device_id_sender,
                device_id_receiver);
            return;
        }

        const std::vector<uint32_t>& ct_args = {
            channel_count, static_cast<uint32_t>(sample_count), static_cast<uint32_t>(sample_size)};

        Program program_sender;
        Program program_receiver;

        tt_metal::CreateKernel(
            program_sender,
            "tt_metal/tools/profiler/sync/sync_device_kernel_sender.cpp",
            eth_sender_core,
            tt_metal::EthernetConfig{.noc = tt_metal::NOC::RISCV_0_default, .compile_args = ct_args});

        tt_metal::CreateKernel(
            program_receiver,
            "tt_metal/tools/profiler/sync/sync_device_kernel_receiver.cpp",
            eth_receiver_core,
            tt_metal::EthernetConfig{.noc = tt_metal::NOC::RISCV_0_default, .compile_args = ct_args});

        try {
            tt::tt_metal::detail::CompileProgram(device_sender, program_sender);
            tt::tt_metal::detail::CompileProgram(device_receiver, program_receiver);
        } catch (std::exception& e) {
            log_error(tt::LogMetal, "Failed compile: {}", e.what());
            throw e;
        }
        tt_metal::detail::LaunchProgram(
            device_sender, program_sender, false /* wait_until_cores_done */, true /* force_slow_dispatch */);
        tt_metal::detail::LaunchProgram(
            device_receiver, program_receiver, false /* wait_until_cores_done */, true /* force_slow_dispatch */);

        tt_metal::detail::WaitProgramDone(device_sender, program_sender, false);
        tt_metal::detail::WaitProgramDone(device_receiver, program_receiver, false);

        CoreCoord sender_core = {eth_sender_core.x, eth_sender_core.y};
        std::vector<CoreCoord> sender_cores = {
            device_sender->virtual_core_from_logical_core(sender_core, CoreType::ETH)};

        CoreCoord receiver_core = {eth_receiver_core.x, eth_receiver_core.y};
        std::vector<CoreCoord> receiver_cores = {
            device_receiver->virtual_core_from_logical_core(receiver_core, CoreType::ETH)};

        peekDeviceData(device_sender, sender_cores);
        peekDeviceData(device_receiver, receiver_cores);

        const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
            tt::tt_metal::MetalContext::instance().profiler_state_manager();
        TT_ASSERT(
            profiler_state_manager->device_profiler_map.at(device_id_sender).device_sync_new_markers.size() ==
            profiler_state_manager->device_profiler_map.at(device_id_receiver).device_sync_new_markers.size());

        auto event_receiver =
            profiler_state_manager->device_profiler_map.at(device_id_receiver).device_sync_new_markers.begin();

        for (auto event_sender =
                 profiler_state_manager->device_profiler_map.at(device_id_sender).device_sync_new_markers.begin();
             event_sender !=
             profiler_state_manager->device_profiler_map.at(device_id_sender).device_sync_new_markers.end();
             event_sender++) {
            TT_ASSERT(
                event_receiver !=
                profiler_state_manager->device_profiler_map.at(device_id_receiver).device_sync_markers.end());
            profiler_state_manager->device_device_time_pair.at(device_id_sender)
                .at(device_id_receiver)
                .push_back({event_sender->timestamp, event_receiver->timestamp});
            event_receiver++;
        }
    }
}

void setSyncInfo(
    chip_id_t device_id,
    std::pair<double, int64_t> syncInfo,
    SyncInfo& root_sync_info,
    std::unordered_map<chip_id_t, std::unordered_map<chip_id_t, std::pair<double, int64_t>>>& deviceDeviceSyncInfo,
    const std::string& parentInfo = "") {
    ZoneScoped;
    const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
        tt::tt_metal::MetalContext::instance().profiler_state_manager();
    if (profiler_state_manager->sync_set_devices.find(device_id) == profiler_state_manager->sync_set_devices.end()) {
        profiler_state_manager->sync_set_devices.insert(device_id);
        if (deviceDeviceSyncInfo.find(device_id) != deviceDeviceSyncInfo.end()) {
            std::string parentInfoNew =
                parentInfo + fmt::format("->{}: ({},{})", device_id, syncInfo.second, syncInfo.first);
            for (auto child_device : deviceDeviceSyncInfo.at(device_id)) {
                std::pair<double, int64_t> childSyncInfo = child_device.second;
                childSyncInfo.second *= syncInfo.first;
                childSyncInfo.second += syncInfo.second;
                childSyncInfo.first *= syncInfo.first;
                setSyncInfo(child_device.first, childSyncInfo, root_sync_info, deviceDeviceSyncInfo, parentInfo);
            }
        }
        detail::setShift(device_id, syncInfo.second, syncInfo.first, root_sync_info);
    }
}

void syncAllDevices(chip_id_t host_connected_device) {
    // Check if profiler on host connected device is initilized
    const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
        tt::tt_metal::MetalContext::instance().profiler_state_manager();
    if (profiler_state_manager->device_profiler_map.find(host_connected_device) ==
        profiler_state_manager->device_profiler_map.end()) {
        return;
    }

    if (!tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_sync_enabled()) {
        return;
    }
    // Update device_device_time_pair
    for (const auto& sender : profiler_state_manager->device_device_time_pair) {
        for (const auto& receiver : sender.second) {
            syncDeviceDevice(sender.first, receiver.first);
        }
    }

    // Run linear regression to calculate scale and bias between devices
    // deviceDeviceSyncInfo[dev0][dev1] = {scale, bias} of dev0 over dev1
    std::unordered_map<chip_id_t, std::unordered_map<chip_id_t, std::pair<double, int64_t>>> deviceDeviceSyncInfo;
    for (auto& sender : profiler_state_manager->device_device_time_pair) {
        for (auto& receiver : sender.second) {
            std::vector<std::pair<uint64_t, uint64_t>> timePairs;
            for (int i = 0; i < receiver.second.size(); i += 2) {
                uint64_t senderTime = (receiver.second[i].first + receiver.second[i + 1].first) / 2;
                timePairs.push_back({senderTime, receiver.second[i].second});
            }
            double senderSum = 0;
            double receiverSum = 0;
            double receiverSquareSum = 0;
            double senderReceiverProductSum = 0;

            // Direct computation causes large error because sqaure of clock is very big
            // So apply linear regression on shifted values
            uint64_t senderBase = 0;
            uint64_t receiverBase = 0;

            if (!timePairs.empty()) {
                senderBase = timePairs[0].first;
                receiverBase = timePairs[0].second;
            }
            for (auto& timePair : timePairs) {
                double senderTime = timePair.first - senderBase;
                double receiverTime = timePair.second - receiverBase;

                receiverSum += receiverTime;
                senderSum += senderTime;
                receiverSquareSum += (receiverTime * receiverTime);
                senderReceiverProductSum += (senderTime * receiverTime);
            }

            uint16_t accumulateSampleCount = timePairs.size();

            double freqScale = (senderReceiverProductSum * accumulateSampleCount - senderSum * receiverSum) /
                               (receiverSquareSum * accumulateSampleCount - receiverSum * receiverSum);

            uint64_t shift = ((double)(senderSum - (freqScale * (double)receiverSum)) / accumulateSampleCount) +
                             (senderBase - freqScale * receiverBase);
            deviceDeviceSyncInfo.insert_or_assign(sender.first, (std::unordered_map<chip_id_t, std::pair<double, int64_t>>){});
            deviceDeviceSyncInfo.at(sender.first)
                .insert_or_assign(receiver.first, (std::pair<double, int64_t>){freqScale, shift});

            deviceDeviceSyncInfo.insert_or_assign(receiver.first, (std::unordered_map<chip_id_t, std::pair<double, int64_t>>){});
            deviceDeviceSyncInfo.at(receiver.first)
                .insert_or_assign(sender.first, (std::pair<double, int64_t>){1.0 / freqScale, -1 * shift});
        }
    }

    // Find any sync info from root device
    // Currently, sync info only exists for SYNC_CORE
    SyncInfo root_sync_info;
    for (auto& [core, info] :
         profiler_state_manager->device_profiler_map.at(host_connected_device).device_core_sync_info) {
        root_sync_info = info;
        break;
    }

    // Propagate sync info with DFS through sync tree
    profiler_state_manager->sync_set_devices.clear();
    setSyncInfo(host_connected_device, (std::pair<double, int64_t>){1.0, 0}, root_sync_info, deviceDeviceSyncInfo);
}

void ProfilerSync(ProfilerSyncState state) {
#if defined(TRACY_ENABLE)
    ZoneScoped;
    if (!tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_sync_enabled()) {
        return;
    }
    if (!getDeviceProfilerState()) {
        return;
    }

    TT_ASSERT(
        !tt::DevicePool::instance().is_dispatch_firmware_active(),
        "Profiler sync is not supported with fast dispatch enabled!");

    const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
        tt::tt_metal::MetalContext::instance().profiler_state_manager();
    static chip_id_t first_connected_device_id = -1;
    if (state == ProfilerSyncState::INIT) {
        profiler_state_manager->do_sync_on_close = true;
        auto ethernet_connections = tt::tt_metal::MetalContext::instance().get_cluster().get_ethernet_connections();
        std::unordered_set<chip_id_t> visited_devices = {};
        constexpr int TOTAL_DEVICE_COUNT = 36;
        for (int sender_device_id = 0; sender_device_id < TOTAL_DEVICE_COUNT; sender_device_id++) {
            if (tt::DevicePool::instance().is_device_active(sender_device_id)) {
                auto sender_device = tt::DevicePool::instance().get_active_device(sender_device_id);
                const auto& active_eth_cores = sender_device->get_active_ethernet_cores(false);

                chip_id_t receiver_device_id;
                tt_xy_pair receiver_eth_core;
                bool doSync = true;
                for (auto& sender_eth_core : active_eth_cores) {
                    if (not tt::tt_metal::MetalContext::instance().get_cluster().is_ethernet_link_up(
                            sender_device_id, sender_eth_core)) {
                        continue;
                    }
                    doSync = false;
                    std::tie(receiver_device_id, receiver_eth_core) =
                        sender_device->get_connected_ethernet_core(sender_eth_core);

                    if (visited_devices.find(sender_device_id) == visited_devices.end() or
                        visited_devices.find(receiver_device_id) == visited_devices.end()) {
                        visited_devices.insert(sender_device_id);
                        visited_devices.insert(receiver_device_id);

                        profiler_state_manager->device_device_time_pair.emplace(
                            sender_device_id,
                            (std::unordered_map<chip_id_t, std::vector<std::pair<uint64_t, uint64_t>>>){});
                        profiler_state_manager->device_device_time_pair.at(sender_device_id)
                            .emplace(receiver_device_id, (std::vector<std::pair<uint64_t, uint64_t>>){});
                    }
                }
                if (doSync or first_connected_device_id == -1) {
                    if (first_connected_device_id == -1 and !doSync) {
                        first_connected_device_id = sender_device_id;
                    }
                    syncDeviceHost(sender_device, ProfilerStateManager::SYNC_CORE, true);
                }
            }
        }
        // If at least one sender receiver pair has been found
        if (first_connected_device_id != -1) {
            syncAllDevices(first_connected_device_id);
        }
    }

    if (state == ProfilerSyncState::CLOSE_DEVICE and profiler_state_manager->do_sync_on_close) {
        profiler_state_manager->do_sync_on_close = false;
        for (const auto& synced_with_host_device : profiler_state_manager->device_host_time_pair) {
            auto deviceToSync = tt::DevicePool::instance().get_active_device(synced_with_host_device.first);
            syncDeviceHost(deviceToSync, ProfilerStateManager::SYNC_CORE, false);
        }
        //  If at least one sender receiver pair has been found
        if (first_connected_device_id != -1) {
            syncAllDevices(first_connected_device_id);
        }
    }
#endif
}

void ClearProfilerControlBuffer(IDevice* device) {
#if defined(TRACY_ENABLE)
    std::vector<uint32_t> control_buffer(kernel_profiler::PROFILER_L1_CONTROL_VECTOR_SIZE, 0);
    setControlBuffer(device, control_buffer);
#endif
}

void InitDeviceProfiler(IDevice* device) {
#if defined(TRACY_ENABLE)
    ZoneScoped;

    TracySetCpuTime(TracyGetCpuTime());

    if (getDeviceProfilerState()) {
        static std::atomic<bool> firstInit = true;

        const chip_id_t device_id = device->id();

        const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
            tt::tt_metal::MetalContext::instance().profiler_state_manager();
        if (profiler_state_manager->device_profiler_map.find(device_id) ==
            profiler_state_manager->device_profiler_map.end()) {
            if (firstInit.exchange(false)) {
                profiler_state_manager->device_profiler_map.try_emplace(device_id, device, true);
            } else {
                profiler_state_manager->device_profiler_map.try_emplace(device_id, device, false);
            }
        }

        auto& soc_desc = tt::tt_metal::MetalContext::instance().get_cluster().get_soc_desc(device_id);

        const uint32_t num_cores_per_dram_bank = soc_desc.profiler_ceiled_core_count_perf_dram_bank;
        const uint32_t bank_size_bytes = PROFILER_FULL_HOST_BUFFER_SIZE_PER_RISC *
                                         MetalContext::instance().hal().get_max_processors_per_core() *
                                         num_cores_per_dram_bank;
        TT_ASSERT(bank_size_bytes <= MetalContext::instance().hal().get_dev_size(HalDramMemAddrType::PROFILER));

        const uint32_t num_dram_banks = soc_desc.get_num_dram_views();

        auto& profiler = profiler_state_manager->device_profiler_map.at(device_id);
        profiler.setLastFDReadAsNotDone();
        profiler.profile_buffer_bank_size_bytes = bank_size_bytes;
        profiler.profile_buffer.resize(profiler.profile_buffer_bank_size_bytes * num_dram_banks / sizeof(uint32_t));

        std::vector<uint32_t> control_buffer(kernel_profiler::PROFILER_L1_CONTROL_VECTOR_SIZE, 0);
        control_buffer[kernel_profiler::DRAM_PROFILER_ADDRESS] =
            MetalContext::instance().hal().get_dev_addr(HalDramMemAddrType::PROFILER);
        setControlBuffer(device, control_buffer);

        if (tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_noc_events_enabled()) {
            profiler.dumpRoutingInfo();
            profiler.dumpClusterCoordinates();
        }
    }
#endif
}

bool areAllCoresDispatchCores(IDevice* device, const std::vector<CoreCoord>& virtual_cores) {
    const chip_id_t device_id = device->id();
    const uint8_t device_num_hw_cqs = device->num_hw_cqs();
    const auto& dispatch_core_config = get_dispatch_core_config();
    std::vector<CoreCoord> dispatch_cores;
    for (const CoreCoord& core : tt::get_logical_dispatch_cores(device_id, device_num_hw_cqs, dispatch_core_config)) {
        const CoreCoord virtual_dispatch_core =
            device->virtual_core_from_logical_core(core, dispatch_core_config.get_core_type());
        dispatch_cores.push_back(virtual_dispatch_core);
    }

    for (const CoreCoord& core : virtual_cores) {
        if (std::find(dispatch_cores.begin(), dispatch_cores.end(), core) == dispatch_cores.end()) {
            return false;
        }
    }
    return true;
}

bool skipReadingDeviceProfilerResults(const ProfilerReadState state) {
    return !tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_do_dispatch_cores() &&
           state == ProfilerReadState::ONLY_DISPATCH_CORES;
}

bool onlyProfileDispatchCores(const ProfilerReadState state) {
    return tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_do_dispatch_cores() &&
           state == ProfilerReadState::ONLY_DISPATCH_CORES;
}

void ReadDeviceProfilerResults(
    IDevice* device,
    const std::vector<CoreCoord>& virtual_cores,
    ProfilerReadState state,
    const std::optional<ProfilerOptionalMetadata>& metadata) {
#if defined(TRACY_ENABLE)
    ZoneScoped;

    if (getDeviceProfilerState()) {
        const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
            tt::tt_metal::MetalContext::instance().profiler_state_manager();
        auto profiler_it = profiler_state_manager->device_profiler_map.find(device->id());
        TT_ASSERT(profiler_it != profiler_state_manager->device_profiler_map.end());
        DeviceProfiler& profiler = profiler_it->second;

        if (skipReadingDeviceProfilerResults(state)) {
            return;
        }

        if (onlyProfileDispatchCores(state)) {
            TT_ASSERT(areAllCoresDispatchCores(device, virtual_cores));

            constexpr uint8_t maxLoopCount = 10;
            constexpr uint32_t loopDuration_us = 10000;

            const auto& hal = MetalContext::instance().hal();
            for (const CoreCoord& core : virtual_cores) {
                bool is_core_done = false;

                const HalProgrammableCoreType core_type = tt::llrt::get_core_type(device->id(), core);

                DeviceAddr profiler_msg_addr = hal.get_dev_addr(core_type, HalL1MemAddrType::PROFILER);
                DeviceAddr control_vector_addr =
                    profiler_msg_addr + hal.get_dev_msgs_factory(core_type).offset_of<dev_msgs::profiler_msg_t>(
                                            dev_msgs::profiler_msg_t::Field::control_vector);
                for (int i = 0; i < maxLoopCount; i++) {
                    const std::vector<std::uint32_t> control_buffer =
                        tt::tt_metal::MetalContext::instance().get_cluster().read_core(
                            device->id(), core, control_vector_addr, kernel_profiler::PROFILER_L1_CONTROL_BUFFER_SIZE);
                    if (control_buffer[kernel_profiler::PROFILER_DONE] == 1) {
                        is_core_done = true;
                        break;
                    }
                    std::this_thread::sleep_for(std::chrono::microseconds(loopDuration_us));
                }
                if (!is_core_done) {
                    std::string msg = fmt::format(
                        "Device profiling never finished on device {}, worker core {}, {}",
                        device->id(),
                        core.x,
                        core.y);
                    TracyMessageC(msg.c_str(), msg.size(), tracy::Color::Tomato3);
                    log_warning(tt::LogMetal, "{}", msg);
                }
            }
        }

        TT_FATAL(
            !tt::tt_metal::MetalContext::instance().dprint_server(),
            "Debug print server is running, cannot read device profiler data");

        if (tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_trace_only()) {
            profiler.readResults(device, virtual_cores, state, ProfilerDataBufferSource::DRAM_AND_L1, metadata);
        } else {
            profiler.readResults(device, virtual_cores, state, ProfilerDataBufferSource::DRAM, metadata);
        }
    }
#endif
}

bool dumpDeviceProfilerDataMidRun(const ProfilerReadState state) {
    if (!tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_mid_run_dump()) {
        return false;
    }

    TT_FATAL(
        tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_mid_run_dump() &&
            !tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_trace_only(),
        "Cannot dump data mid-run if only profiling trace runs");

    TT_FATAL(
        tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_mid_run_dump() &&
            !tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_do_dispatch_cores(),
        "Cannot dump data mid-run if profiling dispatch cores");

    return tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_mid_run_dump() &&
           state == ProfilerReadState::NORMAL;
}

void ProcessDeviceProfilerResults(
    IDevice* device,
    const std::vector<CoreCoord>& virtual_cores,
    ProfilerReadState state,
    const std::optional<ProfilerOptionalMetadata>& metadata) {
#if defined(TRACY_ENABLE)
    ZoneScoped;

    if (getDeviceProfilerState()) {
        const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
            tt::tt_metal::MetalContext::instance().profiler_state_manager();

        auto profiler_it = profiler_state_manager->device_profiler_map.find(device->id());
        TT_ASSERT(profiler_it != profiler_state_manager->device_profiler_map.end());
        DeviceProfiler& profiler = profiler_it->second;

        if (skipReadingDeviceProfilerResults(state)) {
            return;
        }

        if (tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_trace_only()) {
            profiler.processResults(device, virtual_cores, state, ProfilerDataBufferSource::DRAM_AND_L1, metadata);
        } else {
            profiler.processResults(device, virtual_cores, state, ProfilerDataBufferSource::DRAM, metadata);
        }

        if (dumpDeviceProfilerDataMidRun(state)) {
            profiler.dumpDeviceResults(/*is_mid_run_dump=*/true);
        }
    }
#endif
}

std::vector<CoreCoord> getVirtualCoresForProfiling(const IDevice* device, const ProfilerReadState state) {
    std::vector<CoreCoord> virtual_cores;

    const chip_id_t device_id = device->id();
    const uint8_t device_num_hw_cqs = device->num_hw_cqs();
    const auto& dispatch_core_config = get_dispatch_core_config();

    if (!onlyProfileDispatchCores(state)) {
        for (const CoreCoord& core :
             tt::get_logical_compute_cores(device_id, device_num_hw_cqs, dispatch_core_config)) {
            const CoreCoord curr_core = device->worker_core_from_logical_core(core);
            virtual_cores.push_back(curr_core);
        }
        for (const CoreCoord& core : device->get_active_ethernet_cores(true)) {
            const CoreCoord curr_core = device->virtual_core_from_logical_core(core, CoreType::ETH);
            virtual_cores.push_back(curr_core);
        }
    }

    if (tt::tt_metal::MetalContext::instance().rtoptions().get_profiler_do_dispatch_cores()) {
        for (const CoreCoord& core :
             tt::get_logical_dispatch_cores(device_id, device_num_hw_cqs, dispatch_core_config)) {
            const CoreCoord curr_core =
                device->virtual_core_from_logical_core(core, dispatch_core_config.get_core_type());
            virtual_cores.push_back(curr_core);
        }
    }

    return virtual_cores;
}

void ReadDeviceProfilerResults(
    IDevice* device, ProfilerReadState state, const std::optional<ProfilerOptionalMetadata>& metadata) {
#if defined(TRACY_ENABLE)
    ZoneScoped;

    if (getDeviceProfilerState()) {
        TT_ASSERT(device->is_initialized());

        const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
            tt::tt_metal::MetalContext::instance().profiler_state_manager();
        auto profiler_it = profiler_state_manager->device_profiler_map.find(device->id());
        TT_ASSERT(profiler_it != profiler_state_manager->device_profiler_map.end());
        DeviceProfiler& profiler = profiler_it->second;

        if (useFastDispatch(device)) {
            if (profiler.isLastFDReadDone() && state == ProfilerReadState::LAST_FD_READ) {
                ZoneScopedN("Skipping! Last FD dispatch is done");
                return;
            } else if (state == ProfilerReadState::LAST_FD_READ) {
                profiler.setLastFDReadAsDone();
            }
        }

        const std::vector<CoreCoord> virtual_cores = getVirtualCoresForProfiling(device, state);
        ReadDeviceProfilerResults(device, virtual_cores, state, metadata);
        ProcessDeviceProfilerResults(device, virtual_cores, state, metadata);
    }
#endif
}

void SetDeviceProfilerDir(const std::string& output_dir) {
#if defined(TRACY_ENABLE)
    const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
        tt::tt_metal::MetalContext::instance().profiler_state_manager();
    for (auto& device_id : profiler_state_manager->device_profiler_map) {
        profiler_state_manager->device_profiler_map.at(device_id.first).setOutputDir(output_dir);
    }
#endif
}

void FreshProfilerDeviceLog() {
#if defined(TRACY_ENABLE)
    const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
        tt::tt_metal::MetalContext::instance().profiler_state_manager();
    for (auto& device_id : profiler_state_manager->device_profiler_map) {
        profiler_state_manager->device_profiler_map.at(device_id.first).freshDeviceLog();
    }
#endif
}

constexpr uint32_t DEVICE_ID_NUM_BITS = 10;
constexpr uint32_t DEVICE_OP_ID_NUM_BITS = 31;

// Given the base (host assigned id) for a program running on multiple devices, generate a unique per-device
// id by coalescing the physical_device id with the program id.
// For ops running on device, the MSB is 0. For host-fallback ops, the MSB is 1. This avoids aliasing.
uint32_t EncodePerDeviceProgramID(uint32_t base_program_id, uint32_t device_id, bool is_host_fallback_op) {
    return (is_host_fallback_op << DEVICE_OP_ID_NUM_BITS) | (base_program_id << DEVICE_ID_NUM_BITS) | device_id;
}

// Decode per device program ID to get encoded values (base program id, device id, and a flag indicating whether
// it's a host-fallback op).
DeviceProgramId DecodePerDeviceProgramID(uint32_t encoded_device_program_id) {
    DeviceProgramId device_program_id;
    device_program_id.device_id = encoded_device_program_id & ((1 << DEVICE_ID_NUM_BITS) - 1);
    device_program_id.base_program_id =
        (encoded_device_program_id & ((uint32_t)(1 << DEVICE_OP_ID_NUM_BITS) - 1)) >> DEVICE_ID_NUM_BITS;
    device_program_id.is_host_fallback_op = encoded_device_program_id >> DEVICE_OP_ID_NUM_BITS;
    return device_program_id;
}

}  // namespace detail

void ReadMeshDeviceProfilerResults(
    distributed::MeshDevice& mesh_device,
    ProfilerReadState state,
    const std::optional<ProfilerOptionalMetadata>& metadata) {
#if defined(TRACY_ENABLE)
    ZoneScoped;

    if (getDeviceProfilerState()) {
        TT_ASSERT(mesh_device.is_initialized());

        const std::unique_ptr<ProfilerStateManager>& profiler_state_manager =
            tt::tt_metal::MetalContext::instance().profiler_state_manager();

        if (useFastDispatch(&mesh_device)) {
            for (IDevice* device : mesh_device.get_devices()) {
                auto profiler_it = profiler_state_manager->device_profiler_map.find(device->id());
                TT_ASSERT(profiler_it != profiler_state_manager->device_profiler_map.end());
                DeviceProfiler& profiler = profiler_it->second;

                if (profiler.isLastFDReadDone() && state == ProfilerReadState::LAST_FD_READ) {
                    ZoneScopedN("Skipping! Last FD dispatch is done");
                    return;
                } else if (state == ProfilerReadState::LAST_FD_READ) {
                    profiler.setLastFDReadAsDone();
                }
            }

            for (uint8_t cq_id = 0; cq_id < mesh_device.num_hw_cqs(); ++cq_id) {
                mesh_device.mesh_command_queue(cq_id).finish();
            }
        }

        for (IDevice* device : mesh_device.get_devices()) {
            const std::vector<CoreCoord> virtual_cores = detail::getVirtualCoresForProfiling(device, state);
            detail::ReadDeviceProfilerResults(device, virtual_cores, state, metadata);
        }

        for (IDevice* device : mesh_device.get_devices()) {
            mesh_device.enqueue_to_thread_pool([device, state, &metadata]() {
                const std::vector<CoreCoord> virtual_cores = detail::getVirtualCoresForProfiling(device, state);
                detail::ProcessDeviceProfilerResults(device, virtual_cores, state, metadata);
            });
        }

        mesh_device.wait_for_thread_pool();
    }
#endif
}

}  // namespace tt_metal

}  // namespace tt
